import pandas as pd
import numpy as np

def fit_df(*df_list):
      codes=set(df_list[0].columns)
      dates=set(df_list[0].index)
      codes_all=set()
      for df in df_list:
            codes=codes&set(df.columns)
            dates=dates&set(df.index)
            codes_all=codes_all|set(df.columns)
      if len(dates)==0:
            raise ValueError('没有相同的日期')
      if len(codes)==0:
            raise ValueError('没有相同的股票')
      new_df_list=[df.loc[dates,codes] for df in df_list]

      if len(codes_all-codes)>0:
            print('以下股票由于数据不完整,被删除:')
            print(codes_all-codes)
      return new_df_list

def fit_nan(*df_list):
      for df in df_list:
            df=df.isnull()
            df.replace({True:np.nan,False:0},inplace=True)
            try:
                  nan_df
            except:
                  nan_df=df
            else:
                  nan_df=nan_df+df
      return nan_df

class TimeDF(object):
      def __init__(self,df):
            df.sort_index(inplace=True)
            self.df=df
      def __add__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
            else:
                  df2=other
            df=df1+df2
            return TimeDF(df)
      def __sub__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
            else:
                  df2=other
            df1,df2=fit_df(self.df,other.df)
            df=df1-df2
            return TimeDF(df)
      def __mul__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
            else:
                  df2=other
            df1,df2=fit_df(self.df,other.df)
            df=df1*df2
            return TimeDF(df)
      def __truediv__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
            else:
                  df2=other
            df2=df2.replace(0,np.nan)
            df=df1/df2
            return TimeDF(df)
      def __and__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df=fit_nan(df1,df2)
            else:
                  raise TypeError('TimeDF不支持和其它类型比较')
            df=(df1&df2)+nan_df
            df.replace({0:False,1:True},inplace=True)
            return TimeDF(df)
      def __or__(self, other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  raise TypeError('TimeDF不支持和其它类型比较')
            df=(df1|df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)
      def __ge__(self, other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  df1=self.df
                  df2=other
                  nan_df = fit_nan(df1)
            df=(df1 >= df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)
      def __le__(self, other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  df1=self.df
                  df2=other
                  nan_df = fit_nan(df1)
            df=(df1 <= df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)      
      def __eq__(self, other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  df1=self.df
                  df2=other
                  nan_df = fit_nan(df1)
            df=(df1 == df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)
      def __gt__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  df1=self.df
                  df2=other
                  nan_df = fit_nan(df1)
            df=(df1 > df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)
      def __lt__(self,other):
            if isinstance(other,TimeDF):
                  df1,df2=fit_df(self.df,other.df)
                  nan_df = fit_nan(df1, df2)
            else:
                  df1=self.df
                  df2=other
                  nan_df = fit_nan(df1)
            df=(df1 < df2)+nan_df
            df.replace({0: False, 1: True}, inplace=True)
            return TimeDF(df)

